#include <stdio.h>
#include <signal.h>
#include <stdlib.h>

#include <chrono>
#include <thread>
#include <atomic>
#include <iostream>

#include <crane/aio/detail/boost_asyncio.hpp>

using CRA_NS::aio::ByteBuffer;
using CRA_NS::Error;
using AsyncIO = CRA_NS::aio::BoostAsio;


#define tlogf(...) do { \
    fprintf(stderr, "[%#lx] ", static_cast<uint64_t>(std::hash<std::thread::id>{}(std::this_thread::get_id())));    \
    fprintf(stderr, __VA_ARGS__); \
    fprintf(stderr, "\n");  \
} while(0)


class NetBenchClient {
public:
    static const int ThreadCount = 4;

public:
    NetBenchClient(const char *host, int port);
    ~NetBenchClient();

    void start();
    void stop();
    
private:
    void run();
    void onConnect(AsyncIO::handle_t ch, const std::string & host, const std::string& ip, int port, Error err);
    void onSend(AsyncIO::handle_t ch, ByteBuffer && buffer);
    void onRecv(AsyncIO::handle_t ch, ByteBuffer && buffer);
private:
    AsyncIO     asyncio_;
    std::thread threads_[ThreadCount];

    std::atomic<int64_t>      bench_count_ {0};
    std::atomic<int64_t>      bench_time_  {0};
}; // NetBenchClient

NetBenchClient::NetBenchClient(const char *host, int port)
{
    asyncio_.connect(host, port, [this](AsyncIO::handle_t ch, const std::string& host, const std::string& ip, int port, Error err){
        this->onConnect(ch, host, ip, port, std::move(err));
    });
}

NetBenchClient::~NetBenchClient()
{
    std::cout<<"count: "<<bench_count_<<" time: "<<bench_time_<<std::endl;
    std::cout<<"throughput: "<<bench_count_ / (bench_time_.load() / 1000LL ) * 1000 <<" kb/sec"<<std::endl;
}

void NetBenchClient::start()
{
    for ( int i = 0; i < ThreadCount; ++i) {
        threads_[i] = std::thread(&NetBenchClient::run, this);
    }
}

void NetBenchClient::stop()
{
    asyncio_.stop();
    for ( int i = 0; i < ThreadCount; ++i) {
        threads_[i].join();
    }
}

void NetBenchClient::run()
{
    asyncio_.run();
}


void NetBenchClient::onConnect(AsyncIO::handle_t ch, const std::string & host, const std::string& ip, int port, Error err)
{
    assert(!err);

    tlogf("connected to %s:%d, channel: %zu", host.c_str(), port, ch);

    ByteBuffer buffer(1024);
    int64_t * p = (int64_t *)buffer.data();
    *p = CRANE_SYSTEM_TIMESTAMP_USEC();     // 缓存前端写入时间戳
    
    asyncio_.send(ch, std::move(buffer), [this, ch](size_t bytes, ByteBuffer && buffer, Error err) {
        assert(!err); 
        this->onSend(ch, std::move(buffer));
    });
}


void NetBenchClient::onSend(AsyncIO::handle_t ch, ByteBuffer && buffer)
{
    int64_t * p = (int64_t *)buffer.data();
    *p = 0; // 重置头上的时间戳
    buffer.flip();

//    tlogf("sent to channel: %zu",  ch);

    asyncio_.receive(ch, std::move(buffer), [this, ch](size_t bytes, ByteBuffer && buffer, Error err){
        assert(!err);
        this->onRecv(ch, std::move(buffer));
    });
}

void NetBenchClient::onRecv(AsyncIO::handle_t ch, ByteBuffer && buffer)
{
    int64_t * p = (int64_t *)buffer.data();
    int64_t now = CRANE_SYSTEM_TIMESTAMP_USEC();

    bench_count_++;
    bench_time_ += (now - *p);

    if (bench_count_ >= 10000)
    { 
        this->~NetBenchClient();
        exit(0);
    }

    *p = now;
    buffer.flip();
    asyncio_.send(ch, std::move(buffer), [this, ch](size_t bytes, ByteBuffer && buffer, Error err){
        assert(!err);
        this->onSend(ch, std::move(buffer));
    });
}

volatile bool g_stop = false;

static void stop(int sig)
{
    g_stop = true;
}

int main(int argc, char **argv)
{
    if ( argc < 3 ) {
        printf("Usage: %s <host> <port>\n", argv[0]);
        return 1;
    }

    signal(SIGINT, stop);
    signal(SIGTERM, stop);

    const char *host = argv[1];
    int port = atoi(argv[2]);
    NetBenchClient client(host, port);
    client.start();
    
    while (!g_stop) {
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }

    client.stop();
    return 0;
}